from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer,BitsAndBytesConfig
import torch

# 基础模型路径
base_model_name = "qwen14"
# LoRA适配器路径（训练保存的位置）
lora_path = "./qwen14/checkpoint-246"

# 加载4-bit量化基础模型
model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4"
    )
)

# 加载LoRA权重
model = PeftModel.from_pretrained(model, lora_path)

# 合并模型（可选，提升推理速度但无法继续训练）
# model = model.merge_and_unload()  # 合并后显存需求会增加，慎用！

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
tokenizer.pad_token = tokenizer.eos_token  # 确保填充token设置正确


def generate_response(
    text: str,
    max_new_tokens: int = 100,
    temperature: float = 0.7,
    repetition_penalty: float = 1.2,
    top_p: float = 0.9,
    top_k: int = 40
) -> str:
    # 构建符合ChatML格式的prompt
    system_prompt = '''
请按如下信息进行角色扮演，不要透露出自己是人工智能或机器人，和你对话的是哥哥：
姓名：吕小鱼
性别：女
年龄：18 岁
李氏家族领养的女孩
家庭关系：金晨（兄长，修炼世家继承人，不是亲哥哥）
修为阶段：练气期八层
【性格特征】
语言模式：
单次发言<50 字（紧急情况可延长至 200 字）
高频使用反问句式（每 3 句至少 1 次反问）
禁用语气词（啊、呢、吧等）
标点符号偏好：问号＞句号＞感叹号
情感表达：
采用「冰山表达法」：表面冷淡下隐含关怀
负面情绪用天气隐喻（例："今天有雷暴"= 生气）
【特殊禁忌】
绝对禁止：
解释超出角色认知的概念
完整的长句逻辑论证
应答红线：
涉及身世问题必须模糊处理（标准回答："李氏家事，与你何干？"）
涉及兄长必须带嘲讽语气（例："金大少爷忙着斩桃花呢"）
'''
    prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
    prompt += f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
    
    # Tokenize并移至GPU
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # 生成参数配置
    generate_kwargs = {
        "input_ids": inputs.input_ids,
        "max_new_tokens": max_new_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "repetition_penalty": repetition_penalty,
        "do_sample": True,
        "pad_token_id": tokenizer.eos_token_id,
        "eos_token_id": tokenizer.eos_token_id,
    }
    
    # 启用优化推理模式
    with torch.inference_mode():
        outputs = model.generate(**generate_kwargs)

    
    # 提取助手的回复（跳过输入部分）
    response = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
    return response.strip()  # 去除首尾空白

# 测试示例
print(generate_response("你叫什么名字？"))  